package com.thiswhyme.republish.rabbitmq.core;

import cn.hutool.core.lang.TypeReference;
import cn.hutool.json.JSONUtil;
import com.rabbitmq.client.AMQP;
import com.rabbitmq.client.Channel;
import com.thiswhyme.common.expection.BusinessException;
import com.thiswhyme.republish.rabbitmq.config.BindingPlusProperties;
import com.thiswhyme.republish.rabbitmq.config.BindingServicePlusProperties;
import com.thiswhyme.republish.rabbitmq.config.ExecutorConfig;
import com.thiswhyme.republish.rabbitmq.core.domain.MessageInfo;
import com.thiswhyme.republish.rabbitmq.core.domain.MessageType;
import com.thiswhyme.republish.rabbitmq.core.domain.PublishMqHeaders;
import com.thiswhyme.republish.rabbitmq.core.domain.PublishStatus;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.amqp.core.MessageProperties;
import org.springframework.amqp.rabbit.connection.Connection;
import org.springframework.amqp.rabbit.connection.ConnectionFactory;
import org.springframework.amqp.support.AmqpHeaders;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.autoconfigure.jms.JmsProperties;
import org.springframework.context.support.ApplicationObjectSupport;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.support.MessageBuilder;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * @description: 消息重发器
 * @author: thiswhyme
 * @date: 2021/3/18
 */
public class RepublishRunner extends ApplicationObjectSupport implements CommandLineRunner {

    public static final Logger log = LoggerFactory.getLogger(RepublishRunner.class);

    public static final int BATCH_SIZE = 100;

    private ScheduledThreadPoolExecutor toAckScheduler = new ScheduledThreadPoolExecutor(1);

    private ScheduledThreadPoolExecutor toRepublishScheduler = new ScheduledThreadPoolExecutor(1);

    private ThreadPoolTaskExecutor republishExecutor;

    private MessageInfoService messageInfoService;

    private RepublishRedisOperation republishRedisOperation;

    private BindingServicePlusProperties bindingServicePlusProperties;

    private ConnectionFactory connectionFactory;

    public RepublishRunner(ThreadPoolTaskExecutor republishExecutor, MessageInfoService messageInfoService,
                           RepublishRedisOperation republishRedisOperation,
                           BindingServicePlusProperties bindingServicePlusProperties,
                           ConnectionFactory connectionFactory) {
        this.republishExecutor = republishExecutor;
        this.messageInfoService = messageInfoService;
        this.republishRedisOperation = republishRedisOperation;
        this.bindingServicePlusProperties = bindingServicePlusProperties;
        this.connectionFactory = connectionFactory;
    }

    @Override
    public void run(String... args) throws Exception {
        toAckScheduler.scheduleAtFixedRate(() -> {
            try {
                long count = republishRedisOperation.moveToRepublishSet();
                if(count > 0) {
                    log.info("【消息重发器】-准备未ack消息重发，共{}条", count);
                } else {
                    log.trace("【消息重发器】-准备未ack消息重发，共0条");
                }
            } catch (Throwable e) {
                log.error("【消息重发器】-待ack定时器处理异常", e);
            }

        }, 0, 1, TimeUnit.SECONDS);

        toRepublishScheduler.scheduleAtFixedRate(() -> {
            try {
                List<Long> allMessageIds = republishRedisOperation.getAndRemoveFromRepublishSet();
                if(CollectionUtils.isEmpty(allMessageIds)) {
                    if(log.isTraceEnabled()) {
                        log.trace("【消息重发器】-没有待重发的消息，跳过...");
                    }
                } else {
                    List<List<Long>> messageIdBatches = getSubMsgIdList(allMessageIds, BATCH_SIZE);
                    messageIdBatches.forEach(messageIds -> {
                        List<MessageInfo> msgs = messageInfoService.listByIds(messageIds);
                        if(!CollectionUtils.isEmpty(msgs)) {
                            msgs.forEach(this::republish);
                        }
                    });
                    log.info("【消息重发器】-共{}条消息放入重发线程池...", allMessageIds.size());
                }
            } catch (Throwable e) {
                log.error("【消息重发器】-待republish定时器处理异常", e);
            }
        }, 0, 1, TimeUnit.SECONDS);

        //注册线程池关闭钩子
        registerDestroy();
    }

    /**
     * 消息重发
     * @param msg
     */
    private void republish(MessageInfo msg) {
        republishExecutor.execute(() -> {
            try {
                if (MessageType.PRODUCE.getValue() == msg.getType()) {
                    doRepublish(msg);
                } else if (MessageType.CONSUME.getValue() == msg.getType()) {
                    doRetry(msg);
                }
            } catch (Exception e) {
                log.error("【消息重发器】-重发失败，跳过，消息ID：{}", msg.getId(), e);
            }
        });
    }

    /**
     * 重发发送失败的消息
     * @param msg
     * @throws ClassNotFoundException
     */
    private void doRepublish(MessageInfo msg) throws ClassNotFoundException {
        if(null == bindingServicePlusProperties.getBindings()) {
            log.warn("【消息重发器-REPUBLISH】-重发失败，未获取到消息配置，消息ID：{}", msg.getId());
            throw new BusinessException("消息配置不能为空");
        }
        BindingPlusProperties bindingPlusProperties = bindingServicePlusProperties.getBindings().get(msg.getChannelName());
        if (null == bindingPlusProperties) {
            log.warn("【消息重发器-REPUBLISH】-重发失败，未获取到{}通道的消息重发配置，消息ID：{}", msg.getChannelName(), msg.getId());
            throw new BusinessException("消息配置不能为空");
        }
        if(PublishStatus.SEND_TO_ACK.getValue() != msg.getPublishStatus()) {
            log.warn("【消息重发器-REPUBLISH】-跳过重发，消息ID：{}，状态：{}", msg.getId(), msg.getPublishStatus());
        }

        if (bindingPlusProperties.getRepublishTimes() > msg.getRetryTimes()) {
            MessageChannel channel = (MessageChannel) getApplicationContext().getBean(msg.getChannelName());
            //TODO 非json格式的解析问题
            Object payload = JSONUtil.toBean(msg.getContent(), Class.forName(msg.getContentClassType()));
            Map<String, Object> headers = JSONUtil.toBean(msg.getHeaders(), new TypeReference<Map<String, Object>>() {}.getType(), true);
            MessageBuilder builder = MessageBuilder.withPayload(payload);
            builder.copyHeadersIfAbsent(headers).build();
            builder.setHeader(AmqpHeaders.CONTENT_TYPE, null);
            int currentCount = msg.getRetryTimes() + 1;
            builder.setHeader(PublishMqHeaders.HEADER_SEND_RETRY_TIMES, currentCount);
            //消息置入redis待ack的zset
            republishRedisOperation.addToAckZSet(msg.getId(), System.currentTimeMillis()
                    + bindingServicePlusProperties.getAckMillis());

            //更新重发次数
            messageInfoService.updateRetryTimes(msg.getId(), currentCount);

            //重发消息
            channel.send(builder.build());
            if (log.isDebugEnabled()) {
                log.debug("【消息重发器-REPUBLISH】-消息第{}次重发成功，消息ID：{}", currentCount, msg.getId());
            }
        } else {
            messageInfoService.updatePublishStatus(msg.getId(), PublishStatus.FAIL);
            log.warn("【消息重发器-REPUBLISH】-消息已第{}次重发，直接更新状态，消息ID：{}", msg.getRetryTimes(), msg.getId());
        }
    }

    /**
     * 重发消费失败的消息
     * @param msg
     */
    private void doRetry(MessageInfo msg) {
        try (Connection connection = connectionFactory.createConnection();
             Channel channel = connection.createChannel(false)) {
            //组装消息
            Map<String, Object> headers = JSONUtil.toBean(msg.getHeaders(),
                    new TypeReference<Map<String, Object>>() {}.getType(), true);
            headers.put(PublishMqHeaders.HEADER_RECEIVE_RETRY_TIMES,  msg.getRetryTimes() + 1);
            headers.put(PublishMqHeaders.HEADER_RETRY_UNIQUE_ID,  msg.getId());
            headers.put(AmqpHeaders.CONTENT_TYPE, null);
            AMQP.BasicProperties basicProperties = new AMQP.BasicProperties(MessageProperties.CONTENT_TYPE_JSON,
                    null, headers,
                    JmsProperties.DeliveryMode.PERSISTENT.getValue(), 0, null, null,
                    null, null, null, null, null, null, null);
            String exchangeName = (String) headers.get(AmqpHeaders.RECEIVED_EXCHANGE);
            String routingKey = (String) headers.get(AmqpHeaders.RECEIVED_ROUTING_KEY);
            int currentCount = msg.getRetryTimes() + 1;
            log.info("【消息重发器-RETRY】-第{}次重试，exchangeName：{}，routingKey：{}", currentCount, exchangeName,
                    routingKey);

            //更新重发次数
            messageInfoService.updateRetryTimes(msg.getId(), currentCount);

            //重发消息
            channel.basicPublish(exchangeName, routingKey, basicProperties, msg.getContent().getBytes());
        } catch (Exception e) {
            messageInfoService.updatePublishStatus(msg.getId(), PublishStatus.FAIL);
            log.error("【消息重发器-RETRY】-重发消息失败，消息ID：{}", msg.getId(), e);
        }
    }

    /**
     * 将msgIdList分批
     * @param msgIdList 待分批的msgIdList
     * @param len 每批list长度
     * @return
     */
    private List<List<Long>> getSubMsgIdList(List<Long> msgIdList, int len) {
        if (CollectionUtils.isEmpty(msgIdList)) {
            return Collections.emptyList();
        }
        List<List<Long>> resultList = new ArrayList<>();
        for (int i = 0; i < msgIdList.size(); i++) {
            if ( i % len == 0 ) {
                int count = i/len;
                List<Long> subList = msgIdList.stream()
                        .limit((count + 1) * len)
                        .skip(count * len)
                        .collect(Collectors.toList());
                resultList.add(subList);
            }
        }
        return resultList;
    }

    /**
     * 注册线程池关闭钩子
     */
    private void registerDestroy(){
        Runtime.getRuntime().addShutdownHook(new Thread(()->{
            log.info("开始关闭线程池....");
            ExecutorConfig.closeExecutor(toAckScheduler, "待ACK线程池");
            ExecutorConfig.closeExecutor(toRepublishScheduler, "待重发线程池");
            republishExecutor.shutdown();
            log.info("完成关闭线程池....");
        }));
    }
}
